# coding=utf-8
# @author:      ChengJing
# @name:        RNN_model.py
# @datetime:    2021/12/8 15:59
# @software:    PyCharm
# @description:

import torch
import torch.nn as nn

torch.manual_seed(66)
torch.cuda.manual_seed(66)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class MyRNN(nn.Module):
    """
    使用GRU模型构建RNN网络，进行爆管预警和位置的识别
    """
    def __init__(self, in_features, gru_hidden, linear_hidden, out_features):
        """
        Args:
            in_features: 输入数据的特征
            gru_hidden: gru层的隐含层节点数
            linear_hidden: linear层的隐含层节点数
            out_features: 输出数据的特征
        """
        super(MyRNN, self).__init__()
        self.gru = nn.GRU(
            input_size=in_features,
            hidden_size=gru_hidden,
            num_layers=3,
            # dropout=0.5,
            batch_first=True)
        self.classification = nn.Sequential(
            nn.Linear(gru_hidden, linear_hidden),
            nn.ReLU(),
            # nn.Linear(linear_hidden, 2 * linear_hidden),
            # nn.Linear(2 * linear_hidden, linear_hidden),
            nn.Linear(linear_hidden, out_features),
            # nn.Linear(gru_hidden, out_features),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        """
        前向传递函数
        """
        gru_x, _ = self.gru(x)
        lx = gru_x[:, -1, :]
        x = self.classification(lx)
        return x


if __name__ == '__main__':
    model = MyRNN(12, 48, 96, 26)
    from torchsummary import summary
    print(summary(model.cuda(),(60,12)))
